from ultralytics import YOLO
import supervision as sv
import numpy as np

class ObjectDetector:
    """YOLO目标检测模块，负责目标感知与定位"""
    
    def __init__(self, model_path="yolov8n.pt", conf_threshold=0.25, iou_threshold=0.45):
        """初始化目标检测器
        
        Args:
            model_path: YOLO模型路径
            conf_threshold: 置信度阈值
            iou_threshold: IOU阈值
        """
        try:
            self.model = YOLO(model_path)  # 加载YOLO模型
            self.conf_threshold = conf_threshold
            self.iou_threshold = iou_threshold
            self.tracker = sv.ByteTrack()  # 初始化跟踪器
            print(f"YOLO模型 '{model_path}' 加载成功")
        except Exception as e:
            print(f"YOLO模型加载失败: {e}")
            self.model = None
            self.tracker = None
    
    def detect_objects(self, frame):
        """使用YOLO进行目标检测
        
        Args:
            frame: 输入图像（RGB格式）
        
        Returns:
            detections: Supervision Detections对象
            labels: 目标标签列表
        """
        if self.model is None:
            return None, None
        
        # 执行目标检测
        results = self.model(frame, conf=self.conf_threshold, iou=self.iou_threshold, verbose=False)
        
        # 转换为Supervision格式
        detections = sv.Detections.from_ultralytics(results[0])
        
        # 应用跟踪器
        if self.tracker is not None:
            detections = self.tracker.update_with_detections(detections)
        
        # 获取检测到的物体信息
        labels = [
            f"{self.model.names[class_id]} {confidence:.2f} ID:{tracker_id}"
            for class_id, confidence, tracker_id
            in zip(detections.class_id, detections.confidence, detections.tracker_id)
        ]
        
        return detections, labels
    
    def filter_detections_by_class(self, detections, class_names=None):
        """根据类别过滤检测结果
        
        Args:
            detections: Supervision Detections对象
            class_names: 要保留的类别名称列表
        
        Returns:
            过滤后的Detections对象
        """
        if detections is None or len(detections) == 0:
            return detections
        
        if class_names is None or len(class_names) == 0:
            return detections
        
        # 获取类别ID
        class_ids = [self.model.names.index(name) for name in class_names if name in self.model.names]
        
        # 过滤检测结果
        mask = np.isin(detections.class_id, class_ids)
        return detections[mask]
    
    def get_class_counts(self, detections):
        """统计各类别物体的数量
        
        Args:
            detections: Supervision Detections对象
        
        Returns:
            类别计数字典
        """
        if detections is None or len(detections) == 0:
            return {}
        
        class_counts = {}
        for class_id in detections.class_id:
            class_name = self.model.names[class_id]
            class_counts[class_name] = class_counts.get(class_name, 0) + 1
        
        return class_counts
    
    def set_confidence_threshold(self, threshold):
        """设置置信度阈值
        
        Args:
            threshold: 新的置信度阈值 (0-1)
        """
        self.conf_threshold = max(0, min(1, threshold))
    
    def set_iou_threshold(self, threshold):
        """设置IOU阈值
        
        Args:
            threshold: 新的IOU阈值 (0-1)
        """
        self.iou_threshold = max(0, min(1, threshold))

# 注意：不创建全局实例，在实际使用时通过导入类并实例化
# 遵循Java风格的实例化模式，在使用时才创建实例